#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Gamma波实时监测器

本文件专门用于实时监测和显示Gamma波信号。
这是一个轻量级的可视化工具，专注于Gamma波的实时展示。

使用方法：
1. 连接NeuroSky设备
2. 运行此脚本
3. 实时查看Gamma波变化

原理：
1. 使用matplotlib实现实时图表更新
2. 专门监测lowGamma和midGamma信号
3. 提供清晰的波形显示和数值显示
4. 支持数据导出和保存
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'eeg'))

import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.dates import DateFormatter
import numpy as np
from datetime import datetime, timedelta
from collections import deque
import threading
import time

try:
    from eeg_data_logger import EEGDataLogger
except ImportError:
    print("警告: 无法导入eeg_data_logger，将使用模拟数据模式")
    EEGDataLogger = None

class GammaWaveMonitor:
    """Gamma波实时监测器类"""
    
    def __init__(self, port="COM5", baudrate=57600, window_size=30, simulate=False):
        """
        初始化Gamma波监测器
        
        Args:
            port (str): 串口号，默认COM5
            baudrate (int): 波特率，默认57600
            window_size (int): 显示窗口大小（秒），默认30秒
            simulate (bool): 是否使用模拟数据，默认False
        """
        self.port = port
        self.baudrate = baudrate
        self.window_size = window_size
        self.simulate = simulate
        
        # 数据缓存队列
        self.max_points = window_size * 10  # 假设每秒10个数据点
        self.time_data = deque(maxlen=self.max_points)
        self.low_gamma_data = deque(maxlen=self.max_points)
        self.mid_gamma_data = deque(maxlen=self.max_points)
        self.gamma_ratio_data = deque(maxlen=self.max_points)  # Gamma波占比
        
        # 当前数据值
        self.current_low_gamma = 0
        self.current_mid_gamma = 0
        self.current_total_power = 1  # 避免除零
        
        # 图表相关
        self.fig = None
        self.axes = None
        self.lines = {}
        self.ani = None
        
        # 线程控制
        self.data_lock = threading.Lock()
        self.is_monitoring = False
        self.neuropy = None
        
        # 统计信息
        self.gamma_stats = {
            'max_low_gamma': 0,
            'max_mid_gamma': 0,
            'avg_low_gamma': 0,
            'avg_mid_gamma': 0,
            'peak_count': 0
        }
    
    def low_gamma_callback(self, value):
        """低Gamma波回调函数"""
        self.current_low_gamma = value
        self._update_data()
        print(f"低Gamma波: {value}")
    
    def mid_gamma_callback(self, value):
        """中Gamma波回调函数"""
        self.current_mid_gamma = value
        self._update_data()
        print(f"中Gamma波: {value}")
    
    def _update_data(self):
        """更新数据到缓存队列"""
        timestamp = datetime.now()
        gamma_total = self.current_low_gamma + self.current_mid_gamma
        gamma_ratio = gamma_total / self.current_total_power if self.current_total_power > 0 else 0
        
        with self.data_lock:
            self.time_data.append(timestamp)
            self.low_gamma_data.append(self.current_low_gamma)
            self.mid_gamma_data.append(self.current_mid_gamma)
            self.gamma_ratio_data.append(gamma_ratio)
            
            # 更新统计信息
            self._update_stats()
    
    def _update_stats(self):
        """更新统计信息"""
        if len(self.low_gamma_data) > 0:
            self.gamma_stats['max_low_gamma'] = max(self.low_gamma_data)
            self.gamma_stats['max_mid_gamma'] = max(self.mid_gamma_data)
            self.gamma_stats['avg_low_gamma'] = np.mean(list(self.low_gamma_data)[-100:])  # 最近100个点的平均值
            self.gamma_stats['avg_mid_gamma'] = np.mean(list(self.mid_gamma_data)[-100:])
    
    def setup_plots(self):
        """设置图表布局"""
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 创建图表
        self.fig, self.axes = plt.subplots(2, 2, figsize=(14, 8))
        self.fig.suptitle('Gamma波实时监测器', fontsize=16, fontweight='bold', color='red')
        
        # 子图1: Gamma波时间序列
        ax1 = self.axes[0, 0]
        ax1.set_title('Gamma波实时波形', fontweight='bold')
        ax1.set_ylabel('功率值')
        ax1.grid(True, alpha=0.3)
        self.lines['low_gamma'], = ax1.plot([], [], 'g-', label='低Gamma', linewidth=2)
        self.lines['mid_gamma'], = ax1.plot([], [], 'm-', label='中Gamma', linewidth=2)
        self.lines['gamma_total'], = ax1.plot([], [], 'r-', label='总Gamma', linewidth=3, alpha=0.7)
        ax1.legend()
        ax1.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图2: Gamma波比例
        ax2 = self.axes[0, 1]
        ax2.set_title('Gamma波占比', fontweight='bold')
        ax2.set_ylabel('占比')
        ax2.set_ylim(0, 1)
        ax2.grid(True, alpha=0.3)
        self.lines['gamma_ratio'], = ax2.plot([], [], 'orange', label='Gamma占比', linewidth=2)
        ax2.legend()
        ax2.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图3: 实时数值显示
        ax3 = self.axes[1, 0]
        ax3.set_title('实时数值', fontweight='bold')
        ax3.axis('off')  # 隐藏坐标轴
        self.text_display = ax3.text(0.1, 0.5, '', fontsize=12, verticalalignment='center',
                                   bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.8))
        
        # 子图4: Gamma波强度柱状图
        ax4 = self.axes[1, 1]
        ax4.set_title('Gamma波强度对比', fontweight='bold')
        ax4.set_ylabel('功率值')
        
        plt.tight_layout()
        plt.subplots_adjust(top=0.93)
    
    def update_plots(self, frame):
        """更新图表数据"""
        with self.data_lock:
            if len(self.time_data) == 0:
                return list(self.lines.values())
            
            # 获取当前数据
            times = list(self.time_data)
            low_gamma = list(self.low_gamma_data)
            mid_gamma = list(self.mid_gamma_data)
            gamma_ratio = list(self.gamma_ratio_data)
        
        # 计算总Gamma
        gamma_total = [lg + mg for lg, mg in zip(low_gamma, mid_gamma)]
        
        # 更新时间序列图
        self.lines['low_gamma'].set_data(times, low_gamma)
        self.lines['mid_gamma'].set_data(times, mid_gamma)
        self.lines['gamma_total'].set_data(times, gamma_total)
        self.lines['gamma_ratio'].set_data(times, gamma_ratio)
        
        # 动态调整x轴范围
        if len(times) > 0:
            current_time = times[-1]
            start_time = current_time - timedelta(seconds=self.window_size)
            
            self.axes[0, 0].set_xlim(start_time, current_time)
            self.axes[0, 1].set_xlim(start_time, current_time)
        
        # 动态调整y轴范围
        if len(gamma_total) > 10:
            max_gamma = max(gamma_total[-50:]) if gamma_total else 1
            if max_gamma > 0:
                self.axes[0, 0].set_ylim(0, max_gamma * 1.1)
        
        # 更新实时数值显示
        if len(times) > 0:
            current_low = low_gamma[-1] if low_gamma else 0
            current_mid = mid_gamma[-1] if mid_gamma else 0
            current_total = gamma_total[-1] if gamma_total else 0
            current_ratio = gamma_ratio[-1] if gamma_ratio else 0
            
            text_content = f"""当前Gamma波状态:
            
低Gamma: {current_low:.0f}
中Gamma: {current_mid:.0f}
总Gamma: {current_total:.0f}
占比: {current_ratio:.3f}

统计信息:
最大低Gamma: {self.gamma_stats['max_low_gamma']:.0f}
最大中Gamma: {self.gamma_stats['max_mid_gamma']:.0f}
平均低Gamma: {self.gamma_stats['avg_low_gamma']:.1f}
平均中Gamma: {self.gamma_stats['avg_mid_gamma']:.1f}"""
            
            self.text_display.set_text(text_content)
        
        # 更新柱状图
        ax4 = self.axes[1, 1]
        ax4.clear()
        ax4.set_title('Gamma波强度对比', fontweight='bold')
        ax4.set_ylabel('功率值')
        
        if len(times) > 0:
            current_low = low_gamma[-1] if low_gamma else 0
            current_mid = mid_gamma[-1] if mid_gamma else 0
            
            bars = ax4.bar(['低Gamma', '中Gamma'], [current_low, current_mid], 
                          color=['green', 'magenta'], alpha=0.7)
            
            # 在柱状图上显示数值
            for bar, value in zip(bars, [current_low, current_mid]):
                if value > 0:
                    ax4.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(current_low, current_mid)*0.01,
                            f'{int(value)}', ha='center', va='bottom', fontsize=10, fontweight='bold')
        
        return list(self.lines.values())
    
    def start_simulation(self):
        """启动模拟数据模式"""
        def simulate_data():
            while self.is_monitoring:
                # 生成模拟的Gamma波数据
                t = time.time()
                # 使用正弦波加噪声模拟Gamma波
                low_gamma = max(0, 1000 + 500 * np.sin(t * 0.1) + np.random.normal(0, 100))
                mid_gamma = max(0, 800 + 400 * np.sin(t * 0.15 + 1) + np.random.normal(0, 80))
                
                self.current_low_gamma = int(low_gamma)
                self.current_mid_gamma = int(mid_gamma)
                self.current_total_power = 10000  # 模拟总功率
                
                self._update_data()
                
                time.sleep(0.1)  # 每100ms更新一次
        
        # 启动模拟数据线程
        self.data_thread = threading.Thread(target=simulate_data)
        self.data_thread.daemon = True
        self.data_thread.start()
        
        print("模拟数据模式已启动")
    
    def start_real_data(self):
        """启动真实数据采集"""
        try:
            if EEGDataLogger is None:
                raise ImportError("无法导入EEGDataLogger")
            
            from NeuroPy3 import NeuroPy
            self.neuropy = NeuroPy(self.port, self.baudrate)
            
            # 设置Gamma波回调函数
            self.neuropy.setCallBack("lowGamma", self.low_gamma_callback)
            self.neuropy.setCallBack("midGamma", self.mid_gamma_callback)
            
            # 启动数据采集
            self.neuropy.start()
            print(f"真实数据采集已启动 - 串口: {self.port}, 波特率: {self.baudrate}")
            
        except Exception as e:
            print(f"真实数据采集启动失败: {e}")
            print("切换到模拟数据模式")
            self.simulate = True
            self.start_simulation()
    
    def start_monitoring(self):
        """开始Gamma波监测"""
        try:
            # 设置图表
            self.setup_plots()
            
            # 启动数据采集
            self.is_monitoring = True
            
            if self.simulate:
                self.start_simulation()
            else:
                self.start_real_data()
            
            # 启动动画
            self.ani = animation.FuncAnimation(
                self.fig, self.update_plots, interval=100,
                blit=False, cache_frame_data=False
            )
            
            print("=== Gamma波实时监测已启动 ===")
            print("关闭图表窗口或按Ctrl+C停止监测")
            
            # 显示图表
            plt.show()
            
        except KeyboardInterrupt:
            print("\n用户中断，停止监测...")
            self.stop_monitoring()
        except Exception as e:
            print(f"监测错误: {e}")
            self.stop_monitoring()
    
    def stop_monitoring(self):
        """停止Gamma波监测"""
        self.is_monitoring = False
        if self.neuropy:
            self.neuropy.stop()
        if self.ani:
            self.ani.event_source.stop()
        plt.close('all')
        print("Gamma波监测已停止")
    
    def save_data(self, filename=None):
        """保存监测数据"""
        if filename is None:
            filename = f"gamma_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
        
        with self.data_lock:
            if len(self.time_data) == 0:
                print("没有数据可保存")
                return
            
            import pandas as pd
            
            data = {
                'timestamp': list(self.time_data),
                'low_gamma': list(self.low_gamma_data),
                'mid_gamma': list(self.mid_gamma_data),
                'gamma_ratio': list(self.gamma_ratio_data)
            }
            
            df = pd.DataFrame(data)
            df.to_csv(filename, index=False)
            print(f"数据已保存到: {filename}")

def main():
    """主函数"""
    print("=== Gamma波实时监测器 ===")
    print("专门用于监测和显示Gamma波信号")
    print()
    
    # 询问用户选择模式
    mode = input("选择运行模式 (1: 真实数据, 2: 模拟数据): ").strip()
    
    simulate = (mode == '2')
    
    if simulate:
        print("使用模拟数据模式")
    else:
        print("使用真实数据模式")
        port = input("请输入串口号 (默认COM5): ").strip() or "COM5"
    
    # 创建监测器实例
    monitor = GammaWaveMonitor(
        port="COM5" if not simulate else None,
        baudrate=57600,
        window_size=30,  # 显示最近30秒的数据
        simulate=simulate
    )
    
    # 启动监测
    monitor.start_monitoring()

if __name__ == "__main__":
    main()